import psutil
import subprocess
import sys
import threading
import time

"""
    请在系统中运行以下命令安装库
    
    --  pip install psutil

"""
# --- 配置 ---
SEPARATOR = "---MEMORY_SEPARATOR---"
POLL_INTERVAL = 0.02  # 20毫秒轮询一次内存

def main():
    """
    通用包装器：接收一个完整的命令，然后执行、监控并报告结果。
    用法: python generic_memory_wrapper.py [要执行的命令...]
    例如: python generic_memory_wrapper.py python my_script.py
    或者: python generic_memory_wrapper.py ./my_program.exe
    """
    if len(sys.argv) < 2:
        print("用法: python generic_memory_wrapper.py [要执行的命令...]", file=sys.stderr)
        sys.exit(1)

    # 将所有后续参数视为要执行的子命令
    command = sys.argv[1:]

    peak_memory_kb = 0
    process = None

    try:
        # 启动子进程，运行我们接收到的命令
        process = subprocess.Popen(
            command,
            stdin=sys.stdin,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding='utf-8',
            errors='replace' # 增加错误处理，防止二进制输出导致解码失败
        )

        # --- 内存监控线程 ---
        monitor_stop_event = threading.Event()

        def monitor_memory():
            nonlocal peak_memory_kb
            try:
                p = psutil.Process(process.pid)
                while not monitor_stop_event.is_set():
                    try:
                        memory_info = p.memory_info()
                        current_rss_kb = memory_info.rss / 1024
                        if current_rss_kb > peak_memory_kb:
                            peak_memory_kb = current_rss_kb
                    except (psutil.NoSuchProcess, psutil.AccessDenied):
                        break
                    time.sleep(POLL_INTERVAL)
            except Exception:
                pass

        monitor_thread = threading.Thread(target=monitor_memory)
        monitor_thread.start()

        stdout_data, stderr_data = process.communicate()

        monitor_stop_event.set()
        monitor_thread.join()

        # --- 输出最终结果 ---
        sys.stdout.write(stdout_data)
        sys.stderr.write(stderr_data)
        print(f"\n {SEPARATOR} {int(peak_memory_kb)}")
        sys.exit(process.returncode)

    except FileNotFoundError:
        print(f"错误: 找不到命令或程序 '{command[0]}'", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"包装器发生未知错误: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()